#!/usr/bin/env python
# vim: set fileencoding=utf-8 :
import struct
import tempfile
import numpy as np

# > big endian
# < little edian
ENDIAN = '>'
HEAD_SIZE = 12
INT4 = 'I'
INT2 = 'H'
FLOAT4 = 'f'
SAMPLE_LEN_BYTE = 4

class HTKParametersNotCompatibleError(Exception):
	pass

class InstanceError(Exception):
	pass

class AccessError(Exception):
	"""
	raised when try to access an frame outside the duration
	"""
	pass

class HTKParameters:
	"""
	Assume big endian
	"""
	def __init__(self,
					 filename=None,
					 samples=None,
					 samp_period=None,
					 parm_kind=0x0009):
		if filename != None and samples == None:
			self._file = open(filename, 'r+b')
			self.n_samples = self.get_int4()[0]
			self.samp_period = self.get_int4()[0]
			self.samp_size = self.get_int2()[0]
			self.parm_kind = self.get_int2()[0]
		elif samples is not None:
			if samp_period == None:
				raise InstanceError('Need samp_period to create a HTKParameters')
			if parm_kind > 0xffff:
				raise InstanceError('Wrong parm_kind')
			if filename == None:
				self._file = tempfile.TemporaryFile()
			else:
				self._file = open(filename, 'w+b')
			self.samples = np.array(samples)
			self.n_samples = len(samples)
			self.samp_period = samp_period
			self.samp_size = SAMPLE_LEN_BYTE * len(samples[0])
			self.parm_kind = parm_kind
			self._write()
		else:
			raise InstanceError('Cannot create instance, no parameters given')
		self.byte_size = self.samp_size / SAMPLE_LEN_BYTE

	def __str__(self):
		s = ('nSamples = %d\n' % self.n_samples +
				'sampPeriod = %d\n' % self.samp_period +
				'sampSize = %d\n' % self.samp_size +
				'parmKind = %d' % self.parm_kind)
		return s

	def d(self, other):
		from dtw import dtw_distance
		return dtw_distance(self, other)

	def _write(self):
		self._file.seek(0,0)
		s = struct.Struct(ENDIAN + INT4)
		self._file.write(s.pack(self.n_samples))
		self._file.write(s.pack(self.samp_period))
		s = struct.Struct(ENDIAN + INT2)
		self._file.write(s.pack(self.samp_size))
		self._file.write(s.pack(self.parm_kind))
		s = struct.Struct(ENDIAN + str(self.n_samples * len(self.samples[0])) + str(FLOAT4))
		self._file.write(s.pack(*self.samples.reshape(np.multiply(*self.samples.shape))))
		self._file.truncate()
		self._file.flush()
		del self.samples

	def get_frames(self, index_list=None):
		"""
		get the frames of index in index_list
		if index_list is None get all frames
		"""
		if index_list == None:
			index_list = range(self.n_samples)
		if isinstance(index_list, int):
			index_list = [index_list]
		frames = []
		for index in index_list:
			if index < 0 or index >= self.n_samples:
				raise AccessError("Cannot access frame '%d' max is '%d'" %
						(index, self.n_samples))
			self._file.seek(self.samp_size * index + HEAD_SIZE, 0)
			frames.append(self.get_floats4(self.byte_size))
		return frames

	def get_frames_concat(self, other_list, index_list=None):
		if index_list == None:
			index_list = range(self.n_samples)
		o_parms_list = []
		for other in other_list:
			if (self.n_samples != other.n_samples):
				raise HTKParametersNotCompatibleError("nSamples is different")
			o_parms_list.append(other.get_frames(index_list))
		m_parms = self.get_frames(index_list)
		n_parms = []
		for i in xrange(len(m_parms)):
			parms = m_parms[i]
			for o_parms in o_parms_list:
				parms += o_parms[i]
			n_parms.append(parms)
		return n_parms

	def get_floats4(self, size=1):
		"""
		read a 4 byte float
		"""
		return struct.unpack(ENDIAN + str(size) + FLOAT4,
				self._file.read(SAMPLE_LEN_BYTE * size))

	def get_int4(self):
		"""
		read a 4 byte integer
		"""
		return struct.unpack(ENDIAN + INT4, self._file.read(4))

	def get_int2(self):
		"""
		read a 2 byte integer
		"""
		return struct.unpack(ENDIAN + INT2, self._file.read(2))

	def close(self):
		self._file.close()


def concatenate_parm_files(base_name_list, extension_list, output_extension='.concat.htk'):
	f = open(base_name_list)
	total = len(f.readlines())
	f.close()
	atual = float(0)
	last = 0
	for base_name in open(base_name_list):
		atual += 1
		if (atual / total) * 100 > last + 3:
			last = int((atual / total) * 100)
			print 'Progress %.1f%%' % ((atual / total) * 100)
		base_name = base_name.strip()
		p_file_list = []
		for extension in extension_list:
			p_file_list.append(HTKParameters(base_name + extension))
		HTKParameters(base_name + output_extension,
				p_file_list[0].get_frames_concat(p_file_list[1:]),
				p_file_list[0].samp_period).close()
		for p_file in p_file_list:
			p_file.close()

if __name__ == '__main__':
	import argparse
	parser = argparse.ArgumentParser('Manipulate HTK parameters file')
	parser.add_argument('-c', '--concatenate', action='store_true',
			help='Concatenate HTK parameters files')
	parser.add_argument('-l', '--base-name-list', required=True,
			help='One base name (name without extension) of a file per line')
	parser.add_argument('-e', '--extensions', nargs='*',
			help='Extensions to concatenate (the parameters are concatenated in this order)',
			default=['.f0frm', '.plp', '.mfcc', '.fbank'])
	parser.add_argument('-o', '--output-extension', default='.concat.htk',
			help='The concatenated file extension')
	args = parser.parse_args()

	if args.concatenate:
		concatenate_parm_files(args.base_name_list,
				args.extensions,
				args.output_extension)
	else:
		print 'Nothing to do...'
